import torch
from torch import nn


class Sampler(nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, logits: torch.Tensor, temperatures: torch.Tensor):
        logits = logits.to(torch.float)  # 转换为 float 类型，避免数值精度问题
        greedy_tokens = logits.argmax(dim=-1) # 在最后一维（词表维度）取最大值的索引（即最可能的token）
        logits.div_(temperatures.unsqueeze(dim=1)) # 用温度调整 logits（除以温度）
        probs = torch.softmax(logits, dim=-1, dtype=torch.float) # 将调整后的 logits 转换为概率分布
        # logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float)
        epsilon = 1e-10  # 防止除零错误
        sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1) + epsilon).argmax(dim=-1) # Gumbel-Max 采样：通过引入噪声实现从概率分布中随机采样
        return torch.where(temperatures == 0, greedy_tokens, sample_tokens)
